import argparse
from ultralytics import YOLO
import os
import sys


def init_args():
    """初始化命令行参数：默认输入为当前目录zhou.png，保留自定义选项"""
    parser = argparse.ArgumentParser(description="YOLOv8 图片/视频识别脚本（默认识别当前目录zhou.png）")
    # 模型选择（默认轻量版yolov8n.pt，平衡速度和精度）
    parser.add_argument("--model", type=str, default="yolov8n.pt",
                        help="YOLOv8 模型权重文件，如 yolov8n.pt/yolov8s.pt（默认：yolov8n.pt）")
    # 输入文件路径（默认当前目录zhou.png，可手动指定其他文件覆盖）
    parser.add_argument("--input", type=str,
                        default=os.path.join(os.path.dirname(sys.argv[0]), "zhou.png"),
                        help=f"输入图片/视频路径（默认：当前脚本目录下的 zhou.png，当前默认路径：{os.path.join(os.path.dirname(sys.argv[0]), 'zhou.png')}）")
    # 输出结果目录（默认 ./yolov8_results，自动创建）
    parser.add_argument("--output", type=str, default="./yolov8_results",
                        help="识别结果保存目录（默认：./yolov8_results）")
    return parser.parse_args()


def check_input_file(input_path):
    """检查输入文件是否存在且格式合法，友好提示错误原因"""
    # 支持的文件格式（图片+视频）
    supported_img = [".jpg", ".jpeg", ".png", ".bmp", ".gif"]
    supported_vid = [".mp4", ".avi", ".mov", ".mkv", ".flv"]
    supported = supported_img + supported_vid
    file_ext = os.path.splitext(input_path)[1].lower()

    # 检查文件是否存在
    if not os.path.exists(input_path):
        raise FileNotFoundError(
            f"❌ 输入文件不存在！\n当前默认文件路径：{input_path}\n请确认 zhou.png 是否在该目录下，或通过 --input 手动指定其他文件路径"
        )

    # 检查文件格式
    if file_ext not in supported:
        raise ValueError(
            f"❌ 不支持的文件格式：{file_ext}\n支持的图片格式：{supported_img}\n支持的视频格式：{supported_vid}"
        )

    # 返回文件类型（用于后续针对性处理）
    return "image" if file_ext in supported_img else "video"


def run_yolov8_detect(args, file_type):
    """加载模型、执行识别、保存结果，输出清晰的进度信息"""
    # 1. 处理输出目录（确保存在，避免保存失败）
    output_dir = os.path.abspath(args.output)
    result_save_dir = os.path.join(output_dir, "detect_results")
    os.makedirs(result_save_dir, exist_ok=True)
    print(f"📁 识别结果将保存到：{result_save_dir}\n")

    # 2. 加载 YOLOv8 模型（首次运行自动下载权重，提示下载进度）
    print(f"🚀 正在加载 YOLOv8 模型：{args.model}")
    try:
        model = YOLO(args.model)
    except Exception as e:
        raise RuntimeError(
            f"❌ 模型加载失败：{str(e)}\n建议检查网络（首次需下载模型）或手动下载权重到 C:\\Users\\Administrator\\.ultralytics\\models\\")

    # 3. 执行识别（针对图片/视频优化参数，避免服务器/GUI报错）
    print(f"\n🔍 开始识别 {file_type} 文件：\n{os.path.abspath(args.input)}")
    results = model(
        source=args.input,
        conf=0.25,  # 置信度阈值（0.25为默认，可调整：值越低识别越灵敏，值越高越严格）
        save=True,  # 自动保存带识别框的结果文件
        project=output_dir,  # 结果根目录
        name="detect_results",  # 结果子目录（在project下）
        show=False,  # 关闭实时显示（Windows无GUI或服务器环境必备，避免报错）
        verbose=False  # 减少冗余日志，只显示关键信息
    )

    # 4. 解析并展示识别结果（统计目标类别和数量，清晰易懂）
    print("\n✅ 识别完成！结果摘要：")
    for idx, result in enumerate(results):
        # 统计每个类别的目标数量
        class_count = {}
        for box in result.boxes:
            class_name = model.names[int(box.cls[0])]  # 获取类别名称（如 person、car）
            class_count[class_name] = class_count.get(class_name, 0) + 1

        # 输出详细信息
        print(f"\n【文件 {idx + 1}】")
        print(f"原始文件：{os.path.basename(args.input)}")
        print(f"识别到的目标：{class_count if class_count else '未识别到任何目标（可尝试降低 --conf 阈值）'}")
        print(f"结果文件：{os.path.join(result_save_dir, os.path.basename(args.input))}")


if __name__ == "__main__":
    try:
        # 解析参数（默认指向当前目录zhou.png）
        args = init_args()
        # 校验输入文件合法性
        file_type = check_input_file(args.input)
        # 执行识别流程
        run_yolov8_detect(args, file_type)
        print("\n🎉 所有操作完成！可前往结果目录查看带识别框的文件。")
    except Exception as e:
        # 友好展示错误信息，方便排查问题
        print(f"\n❌ 程序执行出错：{str(e)}")
        sys.exit(1)
